package com.taobao.android.mnndemo;

import android.Manifest;
import android.content.pm.PackageManager;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.Matrix;
import android.os.AsyncTask;
import android.os.Environment;
import android.os.Process;
import android.support.annotation.NonNull;
import android.support.v4.app.ActivityCompat;
import android.support.v7.app.AppCompatActivity;
import android.os.Bundle;
import android.util.Log;
import android.view.View;
import android.widget.ImageView;
import android.widget.TextView;
import android.widget.Toast;

import com.taobao.android.mnn.MNNForwardType;
import com.taobao.android.mnn.MNNImageProcess;
import com.taobao.android.mnn.MNNNetInstance;
import com.taobao.android.utils.Common;
import com.taobao.android.utils.TxtFileReader;

import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;

public class ImageActivity extends AppCompatActivity implements View.OnClickListener {
    private static final String TAG = "ImageActivity-zxx";
    //    private final String TargetPic = "MobileNet/testcat.jpg";
    private final String TargetPic = "stand.jpg";
    private final String MobileWordsFileName = "MobileNet/synset_words.txt";

    private List<String> mMobileTaiWords;

    private ImageView mImageView;
    private TextView mTextView;
    private TextView mResultText;
    private TextView mTimeText;
    private Bitmap mBitmap;

    private MNNNetInstance mNetInstance;
    private MNNNetInstance.Session mSession;
    private MNNNetInstance.Session.Tensor mInputTensor;

    private final int InputWidth = 192;
    private final int InputHeight = 256;

    private class NetPrepareTask extends AsyncTask<String, Void, String> {
        protected String doInBackground(String... tasks) {
            prepareMobileNet();
            return "success";
        }

        protected void onPostExecute(String result) {
            mTextView.setText("Start MobileNet Inference");
            mTextView.setClickable(true);
            mTextView.post(new Runnable() {
                @Override
                public void run() {
                    mTextView.callOnClick();
                }
            });
        }
    }

    private class ImageProcessResult {
        public String result;
        public float inferenceTimeCost;
    }

    private class ImageProcessTask extends AsyncTask<String, Void, ImageProcessResult> {

        protected ImageProcessResult doInBackground(String... tasks) {
            if (mInputTensor == null) {
                runOnUiThread(new Runnable() {
                    @Override
                    public void run() {
                        Toast.makeText(ImageActivity.this, "模型未加载成功", Toast.LENGTH_SHORT).show();
                    }
                });
                return null;
            }
            /*
             *  convert data to input tensor
             */
            final MNNImageProcess.Config config = new MNNImageProcess.Config();
            // normalization params
            config.mean = new float[]{103.94f, 116.78f, 123.68f};
            config.normal = new float[]{0.017f, 0.017f, 0.017f};
            // input data format
            config.dest = MNNImageProcess.Format.BGR;
            // bitmap transform
            Matrix matrix = new Matrix();
            matrix.postScale(InputWidth / (float) mBitmap.getWidth(), InputHeight / (float) mBitmap.getHeight());
            matrix.invert(matrix);

            MNNImageProcess.convertBitmap(mBitmap, mInputTensor, config, matrix);

            final long startTimestamp = System.nanoTime();
            /**
             * inference
             */
            mSession.run();

            final long endTimestamp = System.nanoTime();
            final float inferenceTimeCost = (endTimestamp - startTimestamp) / 1000000.0f;

            /**
             * also you can use runWithCallback if you concern about some outputs of the middle layers,
             * this method execute inference and also return middle Tensor outputs synchronously.
             */
//                MNNNetInstance.Session.Tensor[] tensors =  mSession.runWithCallback(new String[]{"conv1"});

            /**
             * get output tensor
             */
            MNNNetInstance.Session.Tensor simcc_x = mSession.getOutput("simcc_x");
            MNNNetInstance.Session.Tensor simcc_y = mSession.getOutput("simcc_y");
            float[] result_simcc_x = simcc_x.getFloatData();// get float results
//            Log.d(TAG, "doInBackground() called with: result = [" + Arrays.toString(result) + "]");
            float[] result_simcc_y = simcc_y.getFloatData();// get float results
            Log.d(TAG, "doInBackground: result_simcc_x:" + Arrays.toString(result_simcc_x));
            Log.d(TAG, "doInBackground: result_simcc_y:" + Arrays.toString(result_simcc_y));

            writeStringToFile("result_simcc_x",result_simcc_x);
            writeStringToFile("result_simcc_y", result_simcc_y);
//            Log.d(TAG, "doInBackground: " + Arrays.toString(result));

            // 显示结果
//            List<Map.Entry<Integer, Float>> maybes = new ArrayList<>();
//            for (int i = 0; i < result.length; i++) {
//                float confidence = result[i];
//                if (confidence > 0.01) {
//                    maybes.add(new AbstractMap.SimpleEntry<Integer, Float>(i, confidence));
//                }
//            }

//            Log.i(Common.TAG, "Inference result size=" + result.length + ", maybe=" + maybes.size());

//            Collections.sort(maybes, new Comparator<Map.Entry<Integer, Float>>() {
//                @Override
//                public int compare(Map.Entry<Integer, Float> o1, Map.Entry<Integer, Float> o2) {
//                    if (Math.abs(o1.getValue() - o2.getValue()) <= Float.MIN_NORMAL) {
//                        return 0;
//                    }
//                    return o1.getValue() > o2.getValue() ? -1 : 1;
//                }
//            });

            final StringBuilder sb = new StringBuilder();
            sb.append(result_simcc_x.length);
            sb.append("\n");
            sb.append(result_simcc_y.length);
//            for (Map.Entry<Integer, Float> entry : maybes) {
//                sb.append("物体:");
//                sb.append(mMobileTaiWords.get(entry.getKey()));
//                sb.append(" 自信度:");
//                sb.append(entry.getValue());
//                sb.append("\n");
//            }
            sb.append("\n");
            final ImageProcessResult processResult = new ImageProcessResult();
            processResult.result = sb.toString();
            processResult.inferenceTimeCost = inferenceTimeCost;
            return processResult;
        }

        protected void onPostExecute(ImageProcessResult result) {
            if (result == null) {
                return;
            }
            if (isDestroyed() || isFinishing()) {
                return;
            }

            if (currentTimes >= totalTimes) {
                double average = calculateAverage(totalData);
                mTimeText.setText(String.format("%d次 推理，平均花费:%fms", totalTimes, average));
            } else {
                currentTimes++;
                picIndex++;
                mBitmap = bitmaps.get(picIndex % bitmaps.size());
                mImageView.setImageBitmap(mBitmap);

                mResultText.setText(result.result);
                float inferenceTimeCost = result.inferenceTimeCost;
                totalData.add(inferenceTimeCost);
                mTimeText.setText("第" + currentTimes + "次推理\n" + "cost time：" + inferenceTimeCost + "ms");
                new ImageProcessTask().execute();
            }
        }
    }

    private void writeStringToFile(String name, float[] result_simcc_x) {

        try {
            FileWriter fileWriter = new FileWriter(getExternalCacheDir() + "/" + picIndex + "_" + name + "_result.txt");
            fileWriter.write(Arrays.toString(result_simcc_x));
            fileWriter.flush();
            fileWriter.close();
        } catch (IOException e) {
            Log.e(TAG, "doInBackground: ", e);
        }
    }

    // 方法：计算 List 的平均值
    public static double calculateAverage(List<Float> list) {
        if (list == null || list.isEmpty()) {
            throw new IllegalArgumentException("The list cannot be null or empty");
        }

        // 计算总和
        double sum = 0;
        for (Float num : list) {
            sum += num;
        }

        // 计算并返回平均值
        return sum / list.size();
    }
    private List<Float> totalData = new ArrayList<>();

    private int currentTimes = 1;
    private int totalTimes = 1000;
    private int picIndex;


    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_image);

        mImageView = findViewById(R.id.imageView);
        mTextView = findViewById(R.id.textView);
        mResultText = findViewById(R.id.editText);
        mTimeText = findViewById(R.id.timeText);

        mTextView.setOnClickListener(this);

        decodeBitmap("1.jpg");
        decodeBitmap("2.jpg");
        decodeBitmap("3.jpg");
        decodeBitmap("4.jpg");
        decodeBitmap("5.jpg");
        mBitmap = bitmaps.get(picIndex % bitmaps.size());
        mImageView.setImageBitmap(mBitmap);
//        totalTimes = bitmaps.size();

        mTextView.setText("prepare Mobile Net ...");
        mTextView.setClickable(false);

        if (ActivityCompat.checkSelfPermission(this, Manifest.permission.READ_EXTERNAL_STORAGE) != PackageManager.PERMISSION_GRANTED) {
            ActivityCompat.requestPermissions(this, new String[]{Manifest.permission.READ_EXTERNAL_STORAGE}, 1);
        } else {
            final NetPrepareTask prepareTask = new NetPrepareTask();
            prepareTask.execute("");
        }
    }
    ArrayList<Bitmap> bitmaps = new ArrayList<>();

    private void decodeBitmap(String fileName) {
        Bitmap bitmap;
        AssetManager am = getAssets();
        try {
            final InputStream picStream = am.open(fileName);
            bitmap = BitmapFactory.decodeStream(picStream);
            picStream.close();
            bitmaps.add(bitmap);
        } catch (Throwable t) {
            t.printStackTrace();
        }
    }

    @Override
    public void onRequestPermissionsResult(int requestCode, @NonNull String[] permissions, @NonNull int[] grantResults) {
        super.onRequestPermissionsResult(requestCode, permissions, grantResults);
        if (grantResults.length > 0) {
            final NetPrepareTask prepareTask = new NetPrepareTask();
            prepareTask.execute("");
        }
    }

    private void prepareMobileNet() {

//        String modelPath = getCacheDir() + "/" + "mobilenet_v1.caffe.mnn";
//        String modelPath = Environment.getExternalStorageDirectory() + "/" + "mobilenet_v1.caffe.mnn";
        String modelPath = Environment.getExternalStorageDirectory() + "/" + "end2end.mnn";

        try {
            mMobileTaiWords = TxtFileReader.getUniqueUrls(getBaseContext(), MobileWordsFileName, Integer.MAX_VALUE);
        } catch (Throwable t) {
            t.printStackTrace();
        }

        // create net instance
        mNetInstance = MNNNetInstance.createFromFile(modelPath);

        if (mNetInstance != null) {
            // create session with config
            MNNNetInstance.Config config = new MNNNetInstance.Config();
            config.numThread = 4;// set threads
            config.forwardType = MNNForwardType.FORWARD_OPENGL.type;// set CPU/GPU
            config.forwardType = MNNForwardType.FORWARD_CPU.type;// set CPU/GPU
            /**
             * config middle layer names, if you concern about the output of the middle layer.
             * use session.getOutput("layer name") to get the output of the middle layer.
             */
//        config.saveTensors = new String[]{"conv1"};
            mSession = mNetInstance.createSession(config);

            // get input tensor
            mInputTensor = mSession.getInput("input");
        }
    }

    @Override
    public void onBackPressed() {
        finish();
        android.os.Process.killProcess(Process.myPid());
    }

    @Override
    public void onClick(View view) {
        if (mBitmap == null) {
            return;
        }
//        System.loadLibrary("GLES_mali");
        mResultText.setText("inference result ...");
        ImageProcessTask imageProcessTask = new ImageProcessTask();
        imageProcessTask.execute("");
    }

    @Override
    protected void onDestroy() {

        /**
         * instance release
         */
        if (mNetInstance != null) {
            mNetInstance.release();
            mNetInstance = null;
        }

        super.onDestroy();
    }
}
